import sys

import numpy as np
import random
import pandas as pd
import networkx as nx
from tqdm import tqdm
import argparse
from src.RKGE import train
from keras.models import load_model

import get_string
from src.load_base import load_ratings, data_split, get_records, get_rec, getentity_list, get_trueid


def get_file_line_count(file_path):
    line_count = 0
    with open('entity_list', 'r') as file:
        for line in file:
            line_count += 1
    return line_count

def construct_kg(data_dir, train_records, user_set):

    kg = nx.DiGraph()
    triple_np = pd.read_csv(data_dir + 'kg.txt', delimiter='\t', header=None).values
    # 获得起始点构成的集合
    head_set = set(triple_np[:, 0])
    tail_set = set(triple_np[:, 1])
    entity_list = list(head_set.union(tail_set)) + list(user_set)





    for triple in triple_np:
        head = entity_list.index(triple[0])
        tail = entity_list.index(triple[1])
        relation = triple[2]

        kg.add_edge(head, tail, relation=relation)


    n_relation = len(set(triple_np[:, 2]))

    for user in train_records:
        for item in train_records[user]:
            kg.add_edge(entity_list.index(user), entity_list.index(item), relation=n_relation)
            kg.add_edge(entity_list.index(item), entity_list.index(user), relation=n_relation + 1)
    return kg, entity_list

def get_relation_dict(kg):

    relation_dict = dict()
    for edge in tqdm(kg.edges):

        relation_dict[(edge[0], edge[1])] = kg[edge[0]][edge[1]]['relation']

    return relation_dict


def get_paths(args,zfc):
    np.random.seed(555)
    data_dir = './data/' + args.dataset + '/'
    ratings_np = load_ratings(data_dir)
    entity_num=getentity_list(data_dir)
    train_set, eval_set, test_set = data_split(ratings_np, args.ratio,entity_num)
    item_set = set(ratings_np[:, 1])
    user_set = set(ratings_np[:, 0])
    train_records = get_records(train_set)
    # 获取的都是正样品
    eval_records = get_records(eval_set)
    test_records = get_records(test_set)
    rec = get_rec(train_records, eval_records, test_records, item_set)
    # ??????????????????????
    kg, entity_list = construct_kg(data_dir, train_records, user_set)
    relation_dict = get_relation_dict(kg)

    path_dict = dict()
    new_train_set = []
    # 这段代码是使用networkx库中的all_simple_paths函数找出从user节点到item节点的所有简单路径，并将路径长度为args.path_len + 1的路径筛选出来，最终将这些路径存储到列表paths中。
    #
    # 在具体实现中，all_simple_paths函数的第一个参数kg表示一个有向图对象，第二个参数user表示起始节点，第三个参数item表示终止节点，cutoff参数
    # 表示最大路径长度限制，如果不指定，默认为None，表示不限制最大路径长度。列表推导式的过滤条件是使用if语句实现的，如果一个路径的长度为
    # args.path_len + 1，即从user到item需要经过args.path_len个中间节点，则将该路径添加到列表paths中。
    for pair in tqdm(train_set):
        user = entity_list.index(pair[0])
        item = entity_list.index(pair[1])
        new_train_set.append([user, item, pair[2]])

        paths = [path for path in list(nx.all_simple_paths(kg, user, item, cutoff=args.path_len)) if (len(path) == args.path_len + 1)]

        if len(paths) > 50:
            indices = np.random.choice(len(paths), 50, replace=False)
            paths = [paths[i] for i in indices]

        path_dict[(user, item)] = paths

    new_eval_set = []

    for pair in tqdm(eval_set):
        user = entity_list.index(pair[0])
        item = entity_list.index(pair[1])
        new_eval_set.append([user, item, pair[2]])

        paths = [path for path in list(nx.all_simple_paths(kg, user, item, cutoff=args.path_len)) if
                 (len(path) == args.path_len + 1)]

        if len(paths) > 50:
            indices = np.random.choice(len(paths), 50, replace=False)
            paths = [paths[i] for i in indices]

        path_dict[(user, item)] = paths

    new_test_set = []

    for pair in tqdm(test_set):
        user = entity_list.index(pair[0])
        item = entity_list.index(pair[1])
        new_test_set.append([user, item, pair[2]])

        paths = [path for path in list(nx.all_simple_paths(kg, user, item, cutoff=args.path_len)) if
                 (len(path) == args.path_len + 1)]

        if len(paths) > 50:
            indices = np.random.choice(len(paths), 50, replace=False)
            paths = [paths[i] for i in indices]

        path_dict[(user, item)] = paths

    new_rec = dict()
    for user in tqdm(rec):
        # ??????????????????????????????????????????????????????????????????????????????????????
        new_rec[entity_list.index(user)] = [entity_list.index(i) for i in rec[user]]
        new_user = entity_list.index(user)
        for item in new_rec[new_user]:

            paths = [path for path in list(nx.all_simple_paths(kg, new_user, item, cutoff=args.path_len)) if (len(path) == args.path_len + 1)]

            if len(paths) > 50:
                indices = np.random.choice(len(paths), 50, replace=False)
                paths = [paths[i] for i in indices]

            path_dict[(new_user, item)] = paths

    true_id=get_trueid(data_dir,zfc)
    # 添加代码
    data = np.zeros((1975, 3))  # 创建一个1975行、3列的全0数组

    for i in range(1975):
        data[i, 0] = true_id
        data[i, 1] = i
        data[i, 2] = random.randint(0, 1)
    # 添加代码

    np.save(data_dir+str(args.ratio)+'_'+str(args.path_len)+'_path_dict.npy', path_dict)
    np.save(data_dir+str(args.ratio)+'_relation_dict.npy', relation_dict)
    np.save(data_dir+'_entity_list.npy', entity_list)
    np.save(data_dir + str(args.ratio)+'_train_set.npy', new_train_set)
    np.save(data_dir + str(args.ratio) + '_eval_set.npy', new_eval_set)
    np.save(data_dir + str(args.ratio)+'_test_set.npy', new_test_set)
    np.save(data_dir + str(args.ratio)+'_rec.npy', new_rec)
    # 添加代码
    np.save(data_dir + str(args.ratio) + '_test.npy', data)
    # 添加代码

def get_paths1(args,zfc):
    data_dir = './RKGE-master/data/' + args.dataset + '/'
    # 添加代码
    data = np.zeros((165, 3))  # 创建一个1975行、3列的全0数组

    for i in range(165):
        data[i, 0] = zfc
        data[i, 1] = i
        data[i, 2] = random.randint(0, 1)


    # 添加代码
    np.save(str(data_dir) +str(args.ratio) + '_test.npy', data)
    # 添加代码

def recommend_attractions(district, purpose, type, season, age, transportation):
    # district=str(district)
    # purpose = str(purpose)
    # type = str(type)
    # season = str(season)
    # age=str(age)
    # transportation = str(transportation)
    #
    # params = district + purpose + type + season + age + transportation
    params=district*2880+purpose*720+type*144+season*24+age*4+transportation;
    return params


class Configuration:
    def __init__(self):
        self.dataset = 'zhishitupu'
        self.ratio = 0.2
        self.lr = 5e-3
        self.l2 = 1e-4
        self.batch_size = 1024
        self.device = 'cuda:0'
        self.dim = 20
        self.p = 5

# 创建一个配置对象
config = Configuration()

# # 通过对象访问参数
# print("Dataset:", config.dataset)
# print("Ratio:", config.ratio)
# print("Learning Rate:", config.lr)
# print("L2:", config.l2)
# print("Batch Size:", config.batch_size)
# print("Device:", config.device)
# print("Dimension:", config.dim)
# print("Number of Paths:", config.p)

if __name__ == "__main__":
    district=1
    purpose=1
    type=1
    season=1
    age=1
    transportation=1

    # district = sys.argv[1]
    # purpose = sys.argv[2]
    # type=sys.argv[3]
    # season=sys.argv[4]
    # age=sys.argv[5]
    # transportation=sys.argv[6]
    zfc = recommend_attractions(district, purpose, type, season, age, transportation)
    get_paths1(config, zfc)
    recommended_ids = train(config, True)
    print(','.join(map(str, recommended_ids)))


